import tushare as ts
import pandas as pd
import numpy as np
from datetime import datetime, timedelta
import talib
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib import font_manager
from sklearn.preprocessing import StandardScaler
from sklearn.decomposition import PCA
from sklearn.model_selection import train_test_split
from imblearn.over_sampling import SMOTE
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import (classification_report, confusion_matrix,
                             roc_auc_score, roc_curve, auc,
                             precision_recall_curve, average_precision_score)
import joblib
import warnings

warnings.filterwarnings('ignore')


# 设置中文字体
def set_chinese_font():
    try:
        plt.rcParams['font.sans-serif'] = ['SimHei']
        plt.rcParams['axes.unicode_minus'] = False
        print("已设置中文字体为SimHei")
    except:
        try:
            plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']
            print("已设置中文字体为Arial Unicode MS")
        except:
            try:
                font_path = "msyh.ttc"
                font_prop = font_manager.FontProperties(fname=font_path)
                plt.rcParams['font.family'] = font_prop.get_name()
                print(f"已加载字体文件: {font_path}")
            except Exception as e:
                print(f"无法设置中文字体: {e}")
                print("将使用英文显示")


set_chinese_font()


# 初始化tushare
def init_tushare():
    pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')
    return pro


# 获取股票数据
def get_stock_data(pro, stock_list, start_date, end_date):
    all_data = pd.DataFrame()
    for stock in stock_list:
        try:
            df = pro.daily(ts_code=stock, start_date=start_date, end_date=end_date)
            df = df.sort_values('trade_date')
            all_data = pd.concat([all_data, df], ignore_index=True)
            print(f"已获取 {stock} 的数据，共 {len(df)} 条记录")
        except Exception as e:
            print(f"获取 {stock} 数据时出错: {e}")
    return all_data


# 计算技术指标
def calculate_technical_indicators(df):
    df = df.sort_values('trade_date')
    numeric_cols = ['close', 'open', 'high', 'low', 'vol']
    for col in numeric_cols:
        df[col] = df[col].astype(float)

    # 移动平均线
    df['MA5'] = df['close'].rolling(window=5).mean()
    df['MA10'] = df['close'].rolling(window=10).mean()
    df['MA20'] = df['close'].rolling(window=20).mean()
    df['MA60'] = df['close'].rolling(window=60).mean()

    # 指数移动平均线
    df['EMA12'] = df['close'].ewm(span=12, adjust=False).mean()
    df['EMA26'] = df['close'].ewm(span=26, adjust=False).mean()

    # MACD
    df['MACD'] = df['EMA12'] - df['EMA26']
    df['MACD_Signal'] = df['MACD'].ewm(span=9, adjust=False).mean()
    df['MACD_Hist'] = df['MACD'] - df['MACD_Signal']

    # RSI
    delta = df['close'].diff()
    gain = (delta.where(delta > 0, 0)).rolling(window=14).mean()
    loss = (-delta.where(delta < 0, 0)).rolling(window=14).mean()
    rs = gain / loss
    df['RSI'] = 100 - (100 / (1 + rs))

    # 布林带
    df['BB_upper'], df['BB_middle'], df['BB_lower'] = talib.BBANDS(
        df['close'], timeperiod=20, nbdevup=2, nbdevdn=2, matype=0)

    # KDJ
    low_min = df['low'].rolling(window=9).min()
    high_max = df['high'].rolling(window=9).max()
    rsv = (df['close'] - low_min) / (high_max - low_min) * 100
    df['K'] = rsv.ewm(com=2).mean()
    df['D'] = df['K'].ewm(com=2).mean()
    df['J'] = 3 * df['K'] - 2 * df['D']

    # 其他指标
    df['OBV'] = (np.sign(df['close'].diff()) * df['vol']).fillna(0).cumsum()
    df['ATR'] = talib.ATR(df['high'], df['low'], df['close'], timeperiod=14)
    df['ADX'] = talib.ADX(df['high'], df['low'], df['close'], timeperiod=14)
    df['WILLR'] = talib.WILLR(df['high'], df['low'], df['close'], timeperiod=14)
    df['CCI'] = talib.CCI(df['high'], df['low'], df['close'], timeperiod=14)
    df['MOM'] = talib.MOM(df['close'], timeperiod=10)

    return df


# 创建标签
def create_labels(df):
    df['future_5d_return'] = df['close'].shift(-5) / df['close'] - 1
    df['label'] = np.where(df['future_5d_return'] > 0.05, 2,  # 上涨
                           np.where(df['future_5d_return'] < -0.05, 0,  # 下跌
                                    1))  # 震荡
    return df


# 主流程 - 数据获取与处理
def data_processing_pipeline():
    pro = init_tushare()
    end_date = datetime.now().strftime('%Y%m%d')
    start_date = (datetime.now() - timedelta(days=730)).strftime('%Y%m%d')
    stock_list = pro.stock_basic(exchange='', list_status='L')['ts_code'].tolist()[:15]

    raw_data = get_stock_data(pro, stock_list, start_date, end_date)
    processed_data = pd.DataFrame()

    for name, group in raw_data.groupby('ts_code'):
        try:
            group = calculate_technical_indicators(group)
            group = create_labels(group)
            processed_data = pd.concat([processed_data, group], ignore_index=True)
            print(f"已完成 {name} 的技术指标计算和标签创建")
        except Exception as e:
            print(f"处理 {name} 时出错: {e}")

    processed_data.to_csv('processed_stock_data.csv', index=False)
    print("处理后的数据已保存至 processed_stock_data.csv")
    return processed_data


def data_cleaning_feature_engineering(processed_data):
    # 数据清洗
    df_clean = processed_data.dropna()
    numeric_cols = df_clean.select_dtypes(include=[np.number]).columns.tolist()
    numeric_cols.remove('label')
    numeric_cols.remove('future_5d_return')

    # 异常值处理
    z_scores = np.abs((df_clean[numeric_cols] - df_clean[numeric_cols].mean()) / df_clean[numeric_cols].std())
    df_clean = df_clean[(z_scores < 3).all(axis=1)]

    # 数据标准化
    scaler = StandardScaler()
    X = df_clean[numeric_cols]
    y = df_clean['label']
    X_scaled = scaler.fit_transform(X)

    # PCA分析
    pca = PCA()
    X_pca = pca.fit_transform(X_scaled)

    plt.figure(figsize=(10, 6))
    plt.plot(np.cumsum(pca.explained_variance_ratio_))
    plt.xlabel('主成分数量')
    plt.ylabel('累计解释方差')
    plt.title('PCA解释方差')
    plt.grid()
    plt.savefig('pca_explained_variance.png', dpi=300, bbox_inches='tight')
    plt.close()

    n_components = np.argmax(np.cumsum(pca.explained_variance_ratio_) >= 0.95) + 1
    pca = PCA(n_components=n_components)
    X_pca = pca.fit_transform(X_scaled)

    # 相关性分析
    plt.figure(figsize=(16, 12))
    sns.heatmap(df_clean[numeric_cols].corr(), annot=False, cmap='coolwarm', center=0)
    plt.title('特征相关性热力图')
    plt.savefig('correlation_heatmap.png', dpi=300, bbox_inches='tight')
    plt.close()

    # 个股画像
    key_features = ['close', 'MA20', 'RSI', 'MACD', 'OBV', 'ATR']
    stock_profile = df_clean.groupby('ts_code')[key_features].mean()

    plt.figure(figsize=(12, 8))
    for i, feature in enumerate(key_features):
        plt.subplot(2, 3, i + 1)
        sns.boxplot(data=stock_profile[feature])
        plt.title(f'{feature}分布')
    plt.tight_layout()
    plt.savefig('stock_profile_boxplot.png', dpi=300, bbox_inches='tight')
    plt.close()

    # 数据平衡
    smote = SMOTE(random_state=42)
    X_res, y_res = smote.fit_resample(X_scaled, y)

    # 保存scaler
    joblib.dump(scaler, 'scaler.pkl')

    return X_res, y_res, numeric_cols


def model_training_evaluation(X_res, y_res, numeric_cols):
    # 数据分割
    X_train, X_test, y_train, y_test = train_test_split(
        X_res, y_res, test_size=0.3, random_state=42, stratify=y_res)

    # 模型训练
    rf = RandomForestClassifier(n_estimators=100, random_state=42)
    rf.fit(X_train, y_train)

    # 模型评估
    y_pred = rf.predict(X_test)
    y_prob = rf.predict_proba(X_test)

    # 分类报告
    print("\n分类报告:")
    print(classification_report(y_test, y_pred, target_names=['下跌', '震荡', '上涨']))

    # 混淆矩阵
    plt.figure(figsize=(8, 6))
    cm = confusion_matrix(y_test, y_pred)
    sns.heatmap(cm, annot=True, fmt='d', cmap='Blues',
                xticklabels=['下跌', '震荡', '上涨'],
                yticklabels=['下跌', '震荡', '上涨'])
    plt.title('混淆矩阵')
    plt.xlabel('预测值')
    plt.ylabel('真实值')
    plt.savefig('confusion_matrix.png', dpi=300, bbox_inches='tight')
    plt.close()

    # ROC曲线
    plt.figure(figsize=(10, 8))
    classes = np.unique(y_res)
    for i in classes:
        if sum(y_test == i) > 0:
            fpr, tpr, _ = roc_curve(y_test == i, y_prob[:, i])
            roc_auc = auc(fpr, tpr)
            class_name = ['下跌', '震荡', '上涨'][i]
            plt.plot(fpr, tpr, label=f'{class_name} (AUC = {roc_auc:.2f})')

    plt.plot([0, 1], [0, 1], 'k--')
    plt.xlabel('假正率')
    plt.ylabel('真正率')
    plt.title('ROC曲线(一对多)')
    plt.legend()
    plt.savefig('roc_curve.png', dpi=300, bbox_inches='tight')
    plt.close()

    # 精确率-召回率曲线
    plt.figure(figsize=(10, 8))
    for i in classes:
        if sum(y_test == i) > 0:
            precision, recall, _ = precision_recall_curve(y_test == i, y_prob[:, i])
            ap = average_precision_score(y_test == i, y_prob[:, i])
            class_name = ['下跌', '震荡', '上涨'][i]
            plt.plot(recall, precision, label=f'{class_name} (AP = {ap:.2f})')

    plt.xlabel('召回率')
    plt.ylabel('精确率')
    plt.title('精确率-召回率曲线')
    plt.legend()
    plt.savefig('precision_recall_curve.png', dpi=300, bbox_inches='tight')
    plt.close()

    # 特征重要性
    feature_importance = pd.DataFrame({
        '特征': numeric_cols,
        '重要性': rf.feature_importances_
    }).sort_values('重要性', ascending=False)

    plt.figure(figsize=(12, 8))
    sns.barplot(x='重要性', y='特征', data=feature_importance.head(15))
    plt.title('前15个重要特征')
    plt.tight_layout()
    plt.savefig('feature_importance.png', dpi=300, bbox_inches='tight')
    plt.close()

    # 保存模型
    joblib.dump(rf, 'stock_prediction_model.pkl')
    print("模型已保存为 stock_prediction_model.pkl")

    return rf


if __name__ == "__main__":
    # 数据获取与预处理
    processed_data = data_processing_pipeline()

    # 数据清洗与特征工程
    X_res, y_res, numeric_cols = data_cleaning_feature_engineering(processed_data)

    # 模型训练与评估
    model = model_training_evaluation(X_res, y_res, numeric_cols)

    print("\n股票数据分析与建模流程已完成!")